package socket

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/chicken-team-outside/chicken_transmission/api_messages"
	"github.com/chicken-team-outside/chicken_transmission/utils"
	"github.com/golang/protobuf/proto"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"
	"io"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

type ackCode uint8

func (a ackCode) Error() string {
	switch a {
	case ErrAckSuccess:
		return "success"
	case ErrAckOffline:
		return "client is offline"
	case ErrAckSendFailed:
		return "failed to send message"
	case ErrAckInvalidClientId:
		return "invalid client id"
	case ErrAckMsgExpired:
		return "message expired"
	case ErrAckQueryError:
		return "query error"
	default:
		return fmt.Sprintf("unknown ack code %d", a)
	}
}

// ClientOnMessageFunc
//
//	用户自定义的消息接收处理函数
//	ctx 当发送者或者与服务器为离线状态时为取消状态
//	header 消息头
//	body 消息体
type ClientOnMessageFunc func(ctx context.Context, header *api_messages.MessageHeader, body []byte)

// ClientStatusChangeHandler
//
//	用于接受客户端状态改变回调函数
//	msgType 消息类型
//	data 响应数据：
//	消息类型为 MessageType_AddClient、MessageType_UpdateClientInfo 时为 *ClientStatusInfo
//	消息类型为 MessageType_RemoveClient 时为 clientId
//	消息类型为 MessageType_UpdateClientStatus 时为 *ClientStatusInfo
//	消息类型为 ListAllClientStatus 时为 []*ClientStatusInfo
type ClientStatusChangeHandler func(msgType api_messages.MessageType, data interface{})

const (
	ErrAckSuccess         = ackCode(api_messages.AckStatus_Success)
	ErrAckOffline         = ackCode(api_messages.AckStatus_Offline)
	ErrAckSendFailed      = ackCode(api_messages.AckStatus_SendFailed)
	ErrAckInvalidClientId = ackCode(api_messages.AckStatus_InvalidClientId)
	ErrAckMsgExpired      = ackCode(api_messages.AckStatus_MsgExpired)
	ErrAckQueryError      = ackCode(api_messages.AckStatus_QueryError)
)

var ErrConnectionClosed = errors.New("connection closed")
var ErrTimeout = errors.New("wait timeout")
var ErrConnectionHasBegun = errors.New("connection has begun")

type WebSocketClient struct {
	conn           atomic.Value
	onMessage      ClientOnMessageFunc
	onStatusChange ClientStatusChangeHandler
	maxMagSize     int
	replyMap       sync.Map
	msgId          uint64
	msgIdLock      sync.Mutex
	connLock       sync.Mutex
	onConnectClose func()
	connected      atomic.Bool
	clientCtxLock  sync.RWMutex
	clientCtx      map[uint64]*clientCtxEntry
	connCtx        context.Context
	connCtxDone    context.CancelCauseFunc
}

func (c *WebSocketClient) GetMsgId() uint64 {
	c.msgIdLock.Lock()
	defer c.msgIdLock.Unlock()
	c.msgId += 1
	return c.msgId
}

func (c *WebSocketClient) GetConn() *WebSocketConnection {
	val := c.conn.Load()
	if val == nil {
		return nil
	}
	return val.(*WebSocketConnection)
}

func (c *WebSocketClient) Disconnect() {
	conn := c.GetConn()
	if conn != nil {
		conn.Close()
	}
}

type ConnectionFailedResponse struct {
	Code int    `json:"code"`
	Msg  string `json:"msg"`
}

// Connect 连接到服务器
//
//	如果已经处于连接中则返回 ErrConnectionHasBegun 错误, 连接失败时将返回错误
func (c *WebSocketClient) Connect(ctx context.Context, address string, authHeader http.Header) error {
	if c.connLock.TryLock() {
		defer c.connLock.Unlock()
		c.connCtx, c.connCtxDone = context.WithCancelCause(ctx)
		conn, response, err := websocket.DefaultDialer.DialContext(ctx, address, authHeader)
		if response != nil && response.Body != nil {
			defer response.Body.Close()
		}
		if err != nil {
			if response != nil && response.Body != nil {
				all, err01 := io.ReadAll(response.Body)
				if all != nil && err01 == nil {
					r := new(ConnectionFailedResponse)
					if json.Unmarshal(all, r) == nil {
						return fmt.Errorf("connection failed[code=%d, msg=%s] %w", r.Code, r.Msg, err)
					}
				}
			}
			return err
		}
		client := NewWebSocket(conn, context.Background(), c.onClose, c.resolveMsg, c.maxMagSize)
		c.Disconnect()
		c.conn.Store(client)
		c.connected.Store(true)
		go client.KeepAlive(10 * time.Second)
		go func() {
			err = client.HandleMessage()
			zap.L().Warn("handle message error", zap.Error(err))
		}()
		<-c.connCtx.Done()
		err = utils.GetCauseFromContext(c.connCtx)
		if err == errSetClientListSuccess {
			return nil
		}
		return err
	} else {
		return ErrConnectionHasBegun
	}
}

func (c *WebSocketClient) IsConnected() bool {
	return c.connected.Load()
}

func (c *WebSocketClient) onClose() {
	c.connected.Store(false)
	if c.onConnectClose != nil {
		c.onConnectClose()
	}
}

func (c *WebSocketClient) GetCurConContext() context.Context {
	conn := c.GetConn()
	if conn == nil {
		ctx, cancel := context.WithCancelCause(context.Background())
		cancel(ErrConnectionClosed)
		return ctx
	}
	return conn.ctx
}

type Message struct {
	Ctx    context.Context
	Header *api_messages.MessageHeader
	Body   []byte
}

type ReplayContext struct {
	msgQueue       *utils.SefaQueue
	ReplyChan      chan *Message
	Ctx            context.Context
	cancelFunc     context.CancelCauseFunc
	ackCtx         context.Context
	ackCancelFunc  context.CancelCauseFunc
	removeCallback func()
}

func newReplayContext(ctx context.Context) (reply *ReplayContext) {
	reply = &ReplayContext{ReplyChan: make(chan *Message, 1), msgQueue: utils.NewSefaQueue()}
	reply.Ctx, reply.cancelFunc = context.WithCancelCause(ctx)
	reply.ackCtx, reply.ackCancelFunc = context.WithCancelCause(reply.Ctx)
	return
}

func (c *ReplayContext) Close() {
	c.cancelFunc(context.Canceled)
	c.removeCallback()
}

func (c *ReplayContext) transferMsg() {
	for c.Ctx.Err() == nil {
		v := c.msgQueue.Pop()
		if v == nil {
			return
		} else {
			select {
			case <-c.Ctx.Done():
				return
			case c.ReplyChan <- v.(*Message):
			}
		}
	}
}

func (c *ReplayContext) handleReply(ctx context.Context, header *api_messages.MessageHeader, body []byte) {
	if header.Type == api_messages.MessageType_Ack {
		ackRsp := &api_messages.AckMessage{}
		err := proto.Unmarshal(body, ackRsp)
		zap.L().Debug("receive ack message to queue", zap.Any("Header", header), zap.Any("ackRsp", ackRsp), zap.Error(err))
		if err != nil {
			c.ackCancelFunc(fmt.Errorf("unmarshal ack message error %w", err))
		} else if ackRsp.Code == api_messages.AckStatus_Success {
			c.ackCancelFunc(context.Canceled)
		} else {
			c.ackCancelFunc(ackCode(ackRsp.Code))
		}
	} else {
		if c.ackCtx.Err() == nil {
			c.ackCancelFunc(context.Canceled)
		}
		zap.L().Debug("push message to queue", zap.Any("Header", header))
		c.msgQueue.Push(&Message{Ctx: ctx, Header: header, Body: body})
	}
}

func (c *ReplayContext) waitAck(ctx context.Context) error {
	select {
	case <-c.ackCtx.Done():
		err := utils.GetCauseFromContext(c.ackCtx)
		if err != nil && errors.Is(err, context.Canceled) {
			return nil
		}
		return err
	case <-c.Ctx.Done():
		return utils.GetCauseFromContext(c.Ctx)
	case <-ctx.Done():
		return utils.GetCauseFromContext(ctx)
	}
}

func (c *WebSocketClient) SendProtobufCustom(ctx context.Context, timeout int64, message *api_messages.MessageHeader, body proto.Message, receiveMsg bool) (reply *ReplayContext, err error) {
	bodyBin, err := proto.Marshal(body)
	if err != nil {
		return nil, err
	}
	message.BodySize = int32(len(bodyBin))
	return c.SendBinaryCustom(ctx, timeout, message, bodyBin, receiveMsg)
}

func (c *WebSocketClient) SendMessage(ctx context.Context, timeout int64, message *api_messages.MessageHeader, body io.Reader, receiveMsg bool) (reply *ReplayContext, err error) {
	conn := c.GetConn()
	if conn == nil {
		return nil, ErrConnectionClosed
	}
	clientCtx := c.getClientCtxEntry(message.ClientId)
	if clientCtx == nil || clientCtx.ctx.Err() != nil {
		return nil, ErrAckOffline
	}
	if timeout > 0 {
		message.ExpiredAt = time.Now().Unix() + timeout
	}
	reply = newReplayContext(ctx)
	withTimeoutCtx, withTimeoutCancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
	defer withTimeoutCancel()
	reply.removeCallback = func() {
		c.replyMap.Delete(message.Id)
	}
	c.replyMap.Store(message.Id, ClientOnMessageFunc(reply.handleReply))
	err = conn.SendMessageWithConnCtx(withTimeoutCtx, clientCtx.ctx, message, body)
	if err == nil {
		err = reply.waitAck(withTimeoutCtx)
	}
	if err != nil || !receiveMsg {
		reply.Close()
		reply = nil
	} else if receiveMsg {
		go reply.transferMsg()
	}
	return
}

// 用于处理proxy推送的状态消息，如其它客户端的状态信息的更新
func (c *WebSocketClient) getSystemMessageHandler(header *api_messages.MessageHeader) ClientOnMessageFunc {
	if header.ClientId == 0 {
		switch header.Type {
		case api_messages.MessageType_AddClient:
			return c.onAddClient
		case api_messages.MessageType_RemoveClient:
			return c.onRemoveClient
		case api_messages.MessageType_UpdateClientInfo:
			return c.onUpdateClientInfo
		case api_messages.MessageType_ListAllClientStatus:
			return c.setClientList
		case api_messages.MessageType_UpdateClientStatus:
			return c.onClientStatusChange
		}
	}
	return nil
}

// 用于处理和分发消息
func (c *WebSocketClient) resolveMsg(ctx context.Context, header *api_messages.MessageHeader, body io.Reader) {
	if ctx.Err() != nil {
		return
	}
	var receiver ClientOnMessageFunc
	if header.Reply != 0 {
		value, ok := c.replyMap.Load(header.Reply)
		if ok && value != nil {
			receiver = value.(ClientOnMessageFunc)
		} else {
			zap.L().Warn("Failed to find reply message receiver", zap.Any("header", header))
			return
		}
	} else {
		if header.ClientId == 0 {
			receiver = c.getSystemMessageHandler(header)
		}
		if receiver == nil {
			receiver = c.onMessage
		}
	}
	if receiver != nil {
		size := int(header.BodySize)
		bodyBytes := make([]byte, size)
		tmp := bodyBytes
		for size > 0 {
			n, err := body.Read(tmp)
			if err != nil {
				zap.L().Warn("Failed to read message body", zap.Error(err), zap.Any("header", header))
				return
			} else {
				tmp = tmp[n:]
				size -= n
			}
		}
		if ctx.Err() == nil {
			receiver(c.getClientCtx(header.ClientId), header, bodyBytes)
		}
	}
}

func (c *WebSocketClient) SendBinaryCustom(ctx context.Context, timeout int64, message *api_messages.MessageHeader, body []byte, receiveMsg bool) (reply *ReplayContext, err error) {
	return c.SendMessage(ctx, timeout, message, bytes.NewReader(body), receiveMsg)
}

func (c *WebSocketClient) SendBinary(ctx context.Context, timeout int64, clientId, replay uint64, t api_messages.MessageType, body []byte, receiveMsg bool) (reply *ReplayContext, err error) {
	return c.SendBinaryCustom(ctx, timeout, &api_messages.MessageHeader{
		Type:     t,
		Reply:    replay,
		ClientId: clientId,
		CreateAt: time.Now().Unix(),
		Id:       c.GetMsgId(),
		BodySize: int32(len(body)),
	}, body, receiveMsg)
}

func (c *WebSocketClient) SendProtobuf(ctx context.Context, timeout int64, clientId, replay uint64, t api_messages.MessageType, content proto.Message, receiveMsg bool) (reply *ReplayContext, err error) {
	var body []byte
	if content != nil {
		body, err = proto.Marshal(content)
		if err != nil {
			return nil, err
		}
	}
	return c.SendBinary(ctx, timeout, clientId, replay, t, body, receiveMsg)
}

func NewWebSocketClient(maxMagSize int, onMessage ClientOnMessageFunc, onStatusChange ClientStatusChangeHandler, onConnectClose func()) *WebSocketClient {
	return &WebSocketClient{
		onMessage:      onMessage,
		onStatusChange: onStatusChange,
		maxMagSize:     maxMagSize,
		msgId:          uint64(time.Now().UnixMilli()),
		onConnectClose: onConnectClose,
		clientCtx:      map[uint64]*clientCtxEntry{},
	}
}
